import torch
from torch import nn
from torch.nn import init
import numpy as np
import sys
sys.path.append("..")
import d2lzh_pytorch as d2l

# 读取数据
batch_size = 256
train_iter, test_iter = d2l.load_data_fashion_mnist()

# 定义和初始化模型
num_inputs = 784
num_outputs = 10

class LinearNet(nn.Module):

    # 构造函数
    def __init__(self, num_inputs, num_outputs):
        # 调用父类init函数
        super(LinearNet, self).__init__()
        self.linear = nn.Linear(num_inputs, num_outputs)

    # 处理模型数据流
    def forward(self, x): #  shape:(batch, 1, 28, 28)
        y = self.linear(x.view(x.shape[0], -1))
        return y

# 创建模型对象
net = LinearNet(num_inputs, num_outputs)

# 初始化模型参数
init.normal_(net.linear.weight, mean=0, std=0.01)
init.constant_(net.linear.bias, val=0)

#  softmax和交叉熵损失函数
loss = nn.CrossEntropyLoss()

#  定义优化算法
optimizer = torch.optim.SGD(net.parameters(), lr=0.1)

#  训练模型
num_epochs = 5
d2l.train_ch3(net, train_iter, test_iter, loss, num_epochs, batch_size, None, None, optimizer)

